from functools import cached_property, reduce
from typing import ClassVar
from typing_extensions import Self
from pydantic import BaseModel
import polars as pl
from pysam import FastaFile  # pylint: disable=E0611
from .sequence import Sequence


class Interval(BaseModel):
    chrom: str
    start: int
    end: int
    name: str = "."

    def distance(self, offset: int) -> int:
        start, end = self.start + 1, self.end
        if start <= offset <= end:
            return 0
        distance1, distance2 = start - offset, offset - end
        return distance1 if abs(distance1) < abs(distance2) else -distance2

    def overlap(self, start, end) -> Self:
        return Interval(chrom=self.chrom, start=max(start, self.start), end=min(end, self.end))

    @cached_property
    def size(self) -> int:
        return max(self.end - self.start, 0)


class Transcript(BaseModel):
    GenomeFasta: ClassVar[FastaFile]
    name: str
    gene: str
    chrom: str
    tx_start: int
    tx_end: int
    cds_start: int
    cds_end: int
    exon_count: int
    exon_starts: list[int]
    exon_ends: list[int]
    strand: str
    level: str

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def intron_starts(self) -> list[int]:
        return self.exon_ends[:-1]

    @cached_property
    def intron_ends(self) -> list[int]:
        return self.exon_starts[1:]

    @cached_property
    def exon_intervals(self) -> list[Interval]:
        numbers = range(1, self.exon_count + 1)[:: 1 if self.strand == "+" else -1]
        intervals = list(map(lambda x: Interval(chrom=self.chrom, start=x[0], end=x[1], name=f"{self.name}:exon{x[2]}"), zip(self.exon_starts, self.exon_ends, numbers)))
        return intervals if self.strand == "+" else intervals[::-1]

    @cached_property
    def exon_sizes(self) -> int:
        return list(map(lambda x: x.size, self.exon_intervals))

    @cached_property
    def intron_intervals(self) -> list[Interval]:
        numbers = range(1, self.exon_count)[:: 1 if self.strand == "+" else -1]
        intervals = list(map(lambda x: Interval(chrom=self.chrom, start=x[0], end=x[1], name=f"{self.name}:intron{x[2]}"), zip(self.intron_starts, self.intron_ends, numbers)))
        return intervals if self.strand == "+" else intervals[::-1]

    @cached_property
    def intron_sizes(self) -> int:
        return list(map(lambda x: x.size, self.intron_intervals))

    @cached_property
    def coding_exon_intervals(self) -> list[Interval | None]:
        return list(map(lambda x: x if x.size else None, map(lambda y: y.overlap(self.cds_start, self.cds_end), self.exon_intervals)))

    @cached_property
    def coding_exon_sizes(self) -> int:
        return list(map(lambda x: x.size if x else 0, self.coding_exon_intervals))

    @cached_property
    def cds_intervals(self) -> list[Interval]:
        intervals = list(filter(None, self.coding_exon_intervals))
        numbers = range(1, len(intervals) + 1)[:: 1 if self.strand == "+" else -1]
        return list(map(lambda x: x[0].model_copy(update={"name": f"{self.name}:intron{x[1]}"}), zip(intervals, numbers)))

    @cached_property
    def cds_sizes(self) -> int:
        return list(map(lambda x: x.size, self.cds_intervals))

    @cached_property
    def cds_size(self) -> int:
        return sum(self.cds_sizes)

    @cached_property
    def cds_count(self) -> int:
        return len(self.cds_intervals)

    @cached_property
    def coding_dna(self) -> Sequence:
        """Fetch the coding sequencing of the transcript"""

        return reduce(
            lambda x, y: x + y,
            map(
                lambda k: k if self.strand == "+" else k.reverse_complement(),
                map(lambda t: Sequence(self.GenomeFasta.fetch(t.chrom, t.start, t.end)), self.cds_intervals),
            ),
        )

    @cached_property
    def is_coding(self) -> bool:
        # Coding transcripts have CDS with non-zero length.
        return self.cds_end - self.cds_start > 0

    @cached_property
    def nmd_cutoff(self) -> float | int:
        """
        Fuction: Nonsense-mediated decay (NMD) classification
        See more information about NMD: https://en.wikipedia.org/wiki/Nonsense-mediated_decay
        NMD classify: not occurring in the 3′ most exon or the 3′-most 50 bp of the penultimate exon
        """
        cds_sizes = self.cds_sizes if self.strand == "+" else self.cds_sizes[::-1]
        if self.gene == "GJB2":  # Hearing Loss Guidelines
            return float("inf")
        if len(cds_sizes) <= 1:
            return float("-inf")
        return sum(cds_sizes[:-1]) - min(50, cds_sizes[-2])

    @cached_property
    def closest_potential_start_codon(self) -> pl.DataFrame:
        """
        Find the closest poteintial in-frame start codon
        找到除第一个密码子外最最可能的起始密码子，所以从第二个密码子开始找
        """
        if self.cds_intervals:
            start_codons = list(filter(lambda x: self.coding_dna[3 : x + 3].upper() == "ATG", range(3, self.cds_size, 3)))
            if start_codons:
                records = pl.concat(list(map(lambda x: pl.DataFrame({"chrom": self.chrom, "pos": range(x.start, x.end)}), self.cds_intervals))).filter(
                    pl.col("pos").le(start_codons[0]) if self.strand == "+" else pl.col("pos").ge(start_codons[0])
                )
                if not records.is_empty():
                    return records.with_columns(pl.lit(self.chrom).alias("chrom"), pl.lit(start_codons[0]).alias("start_codon"))
        return pl.DataFrame(schema={"chrom": str, "pos": int, "start_codon": int})
